# SPDX-FileCopyrightText: © 2024 Tenstorrent Inc.

# SPDX-License-Identifier: Apache-2.0

from typing import Optional, Tuple

import torch
import random
import ttnn

from tests.ttnn.utils_for_testing import (
    check_with_pcc,
    start_measuring_time,
    stop_measuring_time,
    get_per_core_size_and_num_cores,
)
from models.common.utility_functions import torch_random
import pytest

# Override the default timeout in seconds for hang detection.
TIMEOUT = 20
random.seed(0)

# List[Tensor] tensors = [<[1, 100, 14, 14]>, <[1, 100, 14, 14]>],
# int dim = 1
# List[Tensor] tensors = [<[1, 1056, 7, 7]>, <[1, 48, 7, 7]>],
# int dim = 1

parameters = {
    "nightly": {
        "concat_specs": [
            {"dim": 1, "shapes": [[1, 1, 1024], [1, 196, 1024]]},
            {"dim": 1, "shapes": [[1, 1, 1024], [1, 49, 1024]]},
            {"dim": 1, "shapes": [[1, 1, 1280], [1, 1369, 1280]]},
            {"dim": -1, "shapes": [[1, 1, 16, 32], [1, 1, 16, 32]]},
            {"dim": -1, "shapes": [[1, 1, 16], [1, 1, 1]]},
            {"dim": 1, "shapes": [[1, 1, 192], [1, 1344, 192], [1, 100, 192]]},
            {"dim": -1, "shapes": [[1, 1, 1], [1, 1, 1]]},
            {"dim": 1, "shapes": [[1, 1, 224, 224], [1, 1, 224, 224], [1, 1, 224, 224]]},
            {"dim": -1, "shapes": [[1, 1, 7, 32], [1, 1, 7, 32]]},
            {"dim": -1, "shapes": [[1, 1, 768], [1, 1, 512]]},
            {"dim": 1, "shapes": [[1, 1, 768], [1, 196, 768]]},
            {"dim": 1, "shapes": [[1, 1, 768], [1, 49, 768]]},
            #   {'dim': -1, 'shapes': [[1, 1, None], [1, 1, 1]]},
            #   {'dim': -1, 'shapes': [[1, 1, None], [1, 1, 1]]},
            {"dim": 1, "shapes": [[1, 100, 14, 14], [1, 100, 14, 14]]},
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 1056, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                    [1, 48, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [[1, 1056, 7, 7], [1, 48, 7, 7], [1, 48, 7, 7], [1, 48, 7, 7], [1, 48, 7, 7], [1, 48, 7, 7]],
            },
            {"dim": 1, "shapes": [[1, 1056, 7, 7], [1, 48, 7, 7], [1, 48, 7, 7], [1, 48, 7, 7], [1, 48, 7, 7]]},
            {"dim": 1, "shapes": [[1, 1056, 7, 7], [1, 48, 7, 7], [1, 48, 7, 7], [1, 48, 7, 7]]},
            {"dim": 1, "shapes": [[1, 1056, 7, 7], [1, 48, 7, 7], [1, 48, 7, 7]]},
            {"dim": 1, "shapes": [[1, 1056, 7, 7], [1, 48, 7, 7]]},
            {"dim": 1, "shapes": [[1, 1056, 7, 7]]},
            {"dim": 1, "shapes": [[1, 112, 14, 14], [1, 288, 14, 14], [1, 64, 14, 14], [1, 64, 14, 14]]},
            {"dim": 2, "shapes": [[1, 12, 1, 64], [1, 12, 1, 64]]},
            {"dim": 2, "shapes": [[1, 12, 4, 64], [1, 12, 1, 64]]},
            {"dim": -2, "shapes": [[1, 12, 45, 64], [1, 12, 1, 64]]},
            {"dim": 1, "shapes": [[1, 12, 56, 56], [1, 12, 56, 56]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': -2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 12, None, 64], [1, 12, 1, 64]]},
            {"dim": 1, "shapes": [[1, 120, 28, 28], [1, 120, 28, 28]]},
            {"dim": 1, "shapes": [[1, 122400, 4], [1, 30600, 4], [1, 7650, 4], [1, 1989, 4], [1, 567, 4]]},
            {"dim": 1, "shapes": [[1, 122400, 91], [1, 30600, 91], [1, 7650, 91], [1, 1989, 91], [1, 567, 91]]},
            {"dim": 1, "shapes": [[1, 12288, 85], [1, 3072, 85], [1, 768, 85]]},
            {"dim": 1, "shapes": [[1, 128, 112, 112], [1, 128, 112, 112]]},
            {"dim": 1, "shapes": [[1, 128, 14, 14], [1, 256, 14, 14], [1, 64, 14, 14], [1, 64, 14, 14]]},
            {"dim": 1, "shapes": [[1, 128, 28, 28], [1, 128, 28, 28], [1, 64, 28, 28], [1, 128, 28, 28]]},
            {"dim": 1, "shapes": [[1, 128, 28, 28], [1, 128, 28, 28]]},
            {"dim": 1, "shapes": [[1, 128, 28, 28], [1, 19, 28, 28], [1, 38, 28, 28]]},
            {"dim": 1, "shapes": [[1, 128, 28, 28], [1, 192, 28, 28], [1, 96, 28, 28], [1, 64, 28, 28]]},
            {
                "dim": 1,
                "shapes": [
                    [1, 128, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 128, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 128, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 128, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 128, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 128, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 128, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 128, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                    [1, 32, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [[1, 128, 28, 28], [1, 32, 28, 28], [1, 32, 28, 28], [1, 32, 28, 28], [1, 32, 28, 28]],
            },
            {"dim": 1, "shapes": [[1, 128, 28, 28], [1, 32, 28, 28], [1, 32, 28, 28], [1, 32, 28, 28]]},
            {"dim": 1, "shapes": [[1, 128, 28, 28], [1, 32, 28, 28], [1, 32, 28, 28]]},
            {"dim": 1, "shapes": [[1, 128, 28, 28], [1, 32, 28, 28]]},
            {"dim": 1, "shapes": [[1, 128, 28, 28]]},
            {"dim": 1, "shapes": [[1, 128, 64, 64], [1, 128, 64, 64]]},
            {"dim": 1, "shapes": [[1, 128, 64, 64], [1, 256, 64, 64]]},
            #   {'dim': 1,
            #    'shapes': [[1, 128, None, None],
            #               [1, 128, None, None]]},
            {"dim": 1, "shapes": [[1, 1280, 16, 16], [1, 1280, 16, 16]]},
            {"dim": 1, "shapes": [[1, 1280, 16, 16], [1, 640, 16, 16]]},
            {"dim": 1, "shapes": [[1, 1280, 32, 32], [1, 640, 32, 32]]},
            {"dim": 1, "shapes": [[1, 1280, 8, 8], [1, 1280, 8, 8]]},
            {"dim": -1, "shapes": [[1, 14, 14, 192], [1, 14, 14, 192], [1, 14, 14, 192], [1, 14, 14, 192]]},
            {"dim": -1, "shapes": [[1, 14, 14, 256], [1, 14, 14, 256], [1, 14, 14, 256], [1, 14, 14, 256]]},
            {"dim": 1, "shapes": [[1, 14, 56, 56], [1, 14, 56, 56], [1, 14, 56, 56], [1, 14, 56, 56], [1, 68, 56, 56]]},
            {"dim": 1, "shapes": [[1, 14, 56, 56], [1, 24, 56, 56], [1, 40, 56, 56], [1, 64, 56, 56]]},
            {"dim": 1, "shapes": [[1, 14, 56, 56], [1, 24, 56, 56], [1, 64, 56, 56]]},
            {"dim": 1, "shapes": [[1, 14, 56, 56], [1, 40, 56, 56]]},
            {"dim": 1, "shapes": [[1, 14, 56, 56], [1, 64, 56, 56]]},
            {"dim": 2, "shapes": [[1, 16, 1, 64], [1, 16, 1, 64]]},
            {"dim": -1, "shapes": [[1, 16, 16, 192], [1, 16, 16, 192], [1, 16, 16, 192], [1, 16, 16, 192]]},
            {"dim": -1, "shapes": [[1, 16, 16, 256], [1, 16, 16, 256], [1, 16, 16, 256], [1, 16, 16, 256]]},
            {"dim": 1, "shapes": [[1, 16, 28, 28], [1, 128, 28, 28]]},
            {
                "dim": 1,
                "shapes": [
                    [1, 16, 28, 28],
                    [1, 16, 28, 28],
                    [1, 16, 28, 28],
                    [1, 16, 28, 28],
                    [1, 16, 28, 28],
                    [1, 16, 28, 28],
                    [1, 16, 28, 28],
                    [1, 16, 28, 28],
                    [1, 134, 28, 28],
                ],
            },
            {"dim": 1, "shapes": [[1, 16, 28, 28], [1, 28, 28, 28], [1, 128, 28, 28]]},
            {"dim": 1, "shapes": [[1, 16, 28, 28], [1, 28, 28, 28], [1, 46, 28, 28], [1, 128, 28, 28]]},
            {
                "dim": 1,
                "shapes": [[1, 16, 28, 28], [1, 28, 28, 28], [1, 46, 28, 28], [1, 78, 28, 28], [1, 128, 28, 28]],
            },
            {"dim": 1, "shapes": [[1, 16, 28, 28], [1, 28, 28, 28], [1, 78, 28, 28]]},
            {"dim": 1, "shapes": [[1, 16, 28, 28], [1, 46, 28, 28]]},
            {"dim": 1, "shapes": [[1, 16, 28, 28], [1, 78, 28, 28]]},
            {"dim": -2, "shapes": [[1, 16, 5, 64], [1, 16, 1, 64]]},
            {"dim": 2, "shapes": [[1, 16, 59, 64], [1, 16, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 16, None, 64], [1, 16, 1, 64]]},
            #   {'dim': -2,
            #    'shapes': [[1, 16, None, 64], [1, 16, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 16, None, 64], [1, 16, 1, 64]]},
            {"dim": 1, "shapes": [[1, 160, 14, 14], [1, 224, 14, 14], [1, 64, 14, 14], [1, 64, 14, 14]]},
            {"dim": 1, "shapes": [[1, 160, 7, 7], [1, 160, 7, 7], [1, 462, 7, 7]]},
            {"dim": 1, "shapes": [[1, 160, 7, 7], [1, 272, 7, 7], [1, 640, 7, 7]]},
            {"dim": 1, "shapes": [[1, 160, 7, 7], [1, 640, 7, 7]]},
            {"dim": -1, "shapes": [[1, 160], [1, 160]]},
            {"dim": 1, "shapes": [[1, 192, 14, 14], [1, 208, 14, 14], [1, 48, 14, 14], [1, 64, 14, 14]]},
            {
                "dim": 1,
                "shapes": [
                    [1, 192, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 192, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 192, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 192, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 192, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 192, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 192, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 192, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                    [1, 48, 28, 28],
                ],
            },
            {
                "dim": 1,
                "shapes": [[1, 192, 28, 28], [1, 48, 28, 28], [1, 48, 28, 28], [1, 48, 28, 28], [1, 48, 28, 28]],
            },
            {"dim": 1, "shapes": [[1, 192, 28, 28], [1, 48, 28, 28], [1, 48, 28, 28], [1, 48, 28, 28]]},
            {"dim": 1, "shapes": [[1, 192, 28, 28], [1, 48, 28, 28], [1, 48, 28, 28]]},
            {"dim": 1, "shapes": [[1, 192, 28, 28], [1, 48, 28, 28]]},
            {"dim": 1, "shapes": [[1, 192, 28, 28]]},
            {"dim": 1, "shapes": [[1, 192, 35, 35], [1, 192, 35, 35]]},
            {"dim": 1, "shapes": [[1, 192, 8, 8], [1, 320, 8, 8], [1, 1024, 8, 8]]},
            {"dim": -1, "shapes": [[1, 1], [1, 1], [1, 1], [1, 1]]},
            {
                "dim": 1,
                "shapes": [
                    [1, 20, 28, 28],
                    [1, 20, 28, 28],
                    [1, 20, 28, 28],
                    [1, 20, 28, 28],
                    [1, 20, 28, 28],
                    [1, 20, 28, 28],
                    [1, 20, 28, 28],
                    [1, 20, 28, 28],
                    [1, 168, 28, 28],
                ],
            },
            {"dim": 1, "shapes": [[1, 20, 28, 28], [1, 20, 28, 28]]},
            {"dim": 1, "shapes": [[1, 20, 28, 28], [1, 256, 28, 28]]},
            {"dim": 1, "shapes": [[1, 20, 28, 28], [1, 34, 28, 28], [1, 256, 28, 28]]},
            {"dim": 1, "shapes": [[1, 20, 28, 28], [1, 34, 28, 28], [1, 58, 28, 28], [1, 256, 28, 28]]},
            {
                "dim": 1,
                "shapes": [[1, 20, 28, 28], [1, 34, 28, 28], [1, 58, 28, 28], [1, 98, 28, 28], [1, 256, 28, 28]],
            },
            {"dim": 1, "shapes": [[1, 20, 28, 28], [1, 34, 28, 28], [1, 98, 28, 28]]},
            {"dim": 1, "shapes": [[1, 20, 28, 28], [1, 58, 28, 28]]},
            {"dim": 1, "shapes": [[1, 20, 28, 28], [1, 98, 28, 28]]},
            {"dim": 3, "shapes": [[1, 23, 40, 128], [1, 23, 40, 128]]},
            {"dim": 1, "shapes": [[1, 24, 112, 112], [1, 24, 112, 112]]},
            {"dim": 1, "shapes": [[1, 240, 14, 14], [1, 240, 14, 14]]},
            {"dim": 1, "shapes": [[1, 2400, 4], [1, 600, 4], [1, 150, 4], [1, 54, 4], [1, 24, 4], [1, 6, 4]]},
            {"dim": 1, "shapes": [[1, 2400, 91], [1, 600, 91], [1, 150, 91], [1, 54, 91], [1, 24, 91], [1, 6, 91]]},
            {"dim": 1, "shapes": [[1, 256, 128, 128], [1, 256, 128, 128], [1, 256, 128, 128], [1, 256, 128, 128]]},
            {"dim": 1, "shapes": [[1, 256, 14, 14], [1, 256, 14, 14], [1, 128, 14, 14], [1, 256, 14, 14]]},
            {"dim": 1, "shapes": [[1, 256, 14, 14], [1, 256, 14, 14]]},
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 256, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                    [1, 32, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [[1, 256, 14, 14], [1, 32, 14, 14], [1, 32, 14, 14], [1, 32, 14, 14], [1, 32, 14, 14]],
            },
            {"dim": 1, "shapes": [[1, 256, 14, 14], [1, 32, 14, 14], [1, 32, 14, 14], [1, 32, 14, 14]]},
            {"dim": 1, "shapes": [[1, 256, 14, 14], [1, 32, 14, 14], [1, 32, 14, 14]]},
            {"dim": 1, "shapes": [[1, 256, 14, 14], [1, 32, 14, 14]]},
            {"dim": 1, "shapes": [[1, 256, 14, 14], [1, 320, 14, 14], [1, 128, 14, 14], [1, 128, 14, 14]]},
            {"dim": 1, "shapes": [[1, 256, 14, 14]]},
            {"dim": 1, "shapes": [[1, 256, 28, 28], [1, 160, 28, 28], [1, 160, 28, 28], [1, 160, 28, 28]]},
            {"dim": 1, "shapes": [[1, 256, 32, 32], [1, 256, 32, 32]]},
            {"dim": 1, "shapes": [[1, 256, 32, 32], [1, 512, 32, 32]]},
            {"dim": 1, "shapes": [[1, 256, 56, 56], [1, 256, 56, 56]]},
            {"dim": 1, "shapes": [[1, 256, 7, 7], [1, 320, 7, 7], [1, 128, 7, 7], [1, 128, 7, 7]]},
            {"dim": 1, "shapes": [[1, 256, 8, 8], [1, 256, 8, 8]]},
            {"dim": 1, "shapes": [[1, 256, 8, 8], [1, 512, 8, 8], [1, 512, 8, 8], [1, 256, 8, 8]]},
            #   {'dim': 1,
            #    'shapes': [[1, 256, None, None],
            #               [1, 256, None, None]]},
            {"dim": -1, "shapes": [[1, 28, 28, 128], [1, 28, 28, 128], [1, 28, 28, 128], [1, 28, 28, 128]]},
            {"dim": -1, "shapes": [[1, 28, 28, 96], [1, 28, 28, 96], [1, 28, 28, 96], [1, 28, 28, 96]]},
            {"dim": 4, "shapes": [[1, 3, 16, 16, 2], [1, 3, 16, 16, 2], [1, 3, 16, 16, 81]]},
            {"dim": 4, "shapes": [[1, 3, 32, 32, 2], [1, 3, 32, 32, 2], [1, 3, 32, 32, 81]]},
            {"dim": 4, "shapes": [[1, 3, 64, 64, 2], [1, 3, 64, 64, 2], [1, 3, 64, 64, 81]]},
            {"dim": 1, "shapes": [[1, 32, 256, 256], [1, 32, 256, 256]]},
            {"dim": -1, "shapes": [[1, 32, 32, 128], [1, 32, 32, 128], [1, 32, 32, 128], [1, 32, 32, 128]]},
            {"dim": -1, "shapes": [[1, 32, 32, 96], [1, 32, 32, 96], [1, 32, 32, 96], [1, 32, 32, 96]]},
            #   {'dim': 1,
            #    'shapes': [[1, 32, None, None],
            #               [1, 32, None, None]]},
            {"dim": 1, "shapes": [[1, 320, 64, 64], [1, 320, 64, 64]]},
            {"dim": 1, "shapes": [[1, 336, 14, 14], [1, 336, 14, 14]]},
            {"dim": 1, "shapes": [[1, 36, 56, 56], [1, 36, 56, 56]]},
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 384, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                    [1, 48, 14, 14],
                ],
            },
            {
                "dim": 1,
                "shapes": [[1, 384, 14, 14], [1, 48, 14, 14], [1, 48, 14, 14], [1, 48, 14, 14], [1, 48, 14, 14]],
            },
            {"dim": 1, "shapes": [[1, 384, 14, 14], [1, 48, 14, 14], [1, 48, 14, 14], [1, 48, 14, 14]]},
            {"dim": 1, "shapes": [[1, 384, 14, 14], [1, 48, 14, 14], [1, 48, 14, 14]]},
            {"dim": 1, "shapes": [[1, 384, 14, 14], [1, 48, 14, 14]]},
            {"dim": 1, "shapes": [[1, 384, 14, 14]]},
            {"dim": 1, "shapes": [[1, 384, 17, 17], [1, 256, 17, 17], [1, 256, 17, 17], [1, 128, 17, 17]]},
            {"dim": 1, "shapes": [[1, 384, 17, 17], [1, 256, 17, 17], [1, 384, 17, 17]]},
            {"dim": 1, "shapes": [[1, 384, 7, 7], [1, 384, 7, 7], [1, 128, 7, 7], [1, 128, 7, 7]]},
            {"dim": 1, "shapes": [[1, 40, 14, 14], [1, 116, 14, 14]]},
            {"dim": 1, "shapes": [[1, 40, 14, 14], [1, 196, 14, 14]]},
            {"dim": 1, "shapes": [[1, 40, 14, 14], [1, 320, 14, 14]]},
            {
                "dim": 1,
                "shapes": [
                    [1, 40, 14, 14],
                    [1, 40, 14, 14],
                    [1, 40, 14, 14],
                    [1, 40, 14, 14],
                    [1, 40, 14, 14],
                    [1, 40, 14, 14],
                    [1, 40, 14, 14],
                    [1, 40, 14, 14],
                    [1, 334, 14, 14],
                ],
            },
            {"dim": 1, "shapes": [[1, 40, 14, 14], [1, 40, 14, 14]]},
            {
                "dim": 1,
                "shapes": [[1, 40, 14, 14], [1, 68, 14, 14], [1, 116, 14, 14], [1, 196, 14, 14], [1, 320, 14, 14]],
            },
            {"dim": 1, "shapes": [[1, 40, 14, 14], [1, 68, 14, 14], [1, 116, 14, 14], [1, 320, 14, 14]]},
            {"dim": 1, "shapes": [[1, 40, 14, 14], [1, 68, 14, 14], [1, 196, 14, 14]]},
            {"dim": 1, "shapes": [[1, 40, 14, 14], [1, 68, 14, 14], [1, 320, 14, 14]]},
            {"dim": -1, "shapes": [[1, 45], [1, 1]]},
            {"dim": 1, "shapes": [[1, 480, 7, 7], [1, 480, 7, 7]]},
            {"dim": -1, "shapes": [[1, 5, 16, 32], [1, 5, 16, 32]]},
            {"dim": 1, "shapes": [[1, 512, 14, 14], [1, 192, 14, 14], [1, 192, 14, 14], [1, 192, 14, 14]]},
            {"dim": 1, "shapes": [[1, 512, 28, 28], [1, 512, 28, 28]]},
            {
                "dim": 1,
                "shapes": [
                    [1, 512, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 512, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 512, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 512, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 512, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 512, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 512, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 512, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 512, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 512, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 512, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [[1, 512, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7]],
            },
            {"dim": 1, "shapes": [[1, 512, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7]]},
            {"dim": 1, "shapes": [[1, 512, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7]]},
            {"dim": 1, "shapes": [[1, 512, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7]]},
            {"dim": 1, "shapes": [[1, 512, 7, 7], [1, 32, 7, 7]]},
            {"dim": 1, "shapes": [[1, 512, 7, 7], [1, 512, 7, 7], [1, 256, 7, 7]]},
            {"dim": 1, "shapes": [[1, 512, 7, 7]]},
            {"dim": 1, "shapes": [[1, 56, 14, 14], [1, 56, 14, 14]]},
            {"dim": 1, "shapes": [[1, 5776, 4], [1, 2166, 4], [1, 600, 4], [1, 150, 4], [1, 36, 4], [1, 4, 4]]},
            {"dim": 1, "shapes": [[1, 5776, 91], [1, 2166, 91], [1, 600, 91], [1, 150, 91], [1, 36, 91], [1, 4, 91]]},
            {"dim": -1, "shapes": [[1, 59], [1, 1]]},
            {"dim": -1, "shapes": [[1, 5], [1, 1]]},
            {"dim": 2, "shapes": [[1, 6, 1, 64], [1, 6, 1, 64]]},
            {"dim": 2, "shapes": [[1, 6, 16, 64], [1, 6, 1, 64]]},
            #   {'dim': 2,
            #    'shapes': [[1, 6, None, 64], [1, 6, 1, 64]]},
            {"dim": 1, "shapes": [[1, 60, 28, 28], [1, 60, 28, 28]]},
            {"dim": 1, "shapes": [[1, 64, 120, 160], [1, 64, 120, 160]]},
            {"dim": 1, "shapes": [[1, 64, 128, 128], [1, 64, 128, 128]]},
            {"dim": 1, "shapes": [[1, 64, 224, 224], [1, 64, 224, 224]]},
            {"dim": 1, "shapes": [[1, 64, 28, 28], [1, 128, 28, 28], [1, 32, 28, 28], [1, 32, 28, 28]]},
            {"dim": 1, "shapes": [[1, 64, 30, 40], [1, 64, 30, 40]]},
            {"dim": 1, "shapes": [[1, 64, 56, 56], [1, 128, 56, 56], [1, 128, 56, 56], [1, 128, 56, 56]]},
            {
                "dim": 1,
                "shapes": [
                    [1, 64, 56, 56],
                    [1, 32, 56, 56],
                    [1, 32, 56, 56],
                    [1, 32, 56, 56],
                    [1, 32, 56, 56],
                    [1, 32, 56, 56],
                    [1, 32, 56, 56],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 64, 56, 56],
                    [1, 32, 56, 56],
                    [1, 32, 56, 56],
                    [1, 32, 56, 56],
                    [1, 32, 56, 56],
                    [1, 32, 56, 56],
                ],
            },
            {"dim": 1, "shapes": [[1, 64, 56, 56], [1, 32, 56, 56], [1, 32, 56, 56], [1, 32, 56, 56], [1, 32, 56, 56]]},
            {"dim": 1, "shapes": [[1, 64, 56, 56], [1, 32, 56, 56], [1, 32, 56, 56], [1, 32, 56, 56]]},
            {"dim": 1, "shapes": [[1, 64, 56, 56], [1, 32, 56, 56], [1, 32, 56, 56]]},
            {"dim": 1, "shapes": [[1, 64, 56, 56], [1, 32, 56, 56]]},
            {"dim": 1, "shapes": [[1, 64, 56, 56], [1, 64, 56, 56]]},
            {"dim": 1, "shapes": [[1, 64, 56, 56]]},
            {"dim": 1, "shapes": [[1, 64, 60, 80], [1, 64, 60, 80]]},
            {"dim": 1, "shapes": [[1, 64, 73, 73], [1, 96, 73, 73]]},
            #   {'dim': 1,
            #    'shapes': [[1, 64, None, None],
            #               [1, 64, None, None]]},
            {"dim": 1, "shapes": [[1, 640, 32, 32], [1, 320, 32, 32]]},
            {"dim": 1, "shapes": [[1, 640, 32, 32], [1, 640, 32, 32]]},
            {"dim": 1, "shapes": [[1, 640, 64, 64], [1, 320, 64, 64]]},
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [
                    [1, 640, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                    [1, 32, 7, 7],
                ],
            },
            {
                "dim": 1,
                "shapes": [[1, 640, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7]],
            },
            {"dim": 1, "shapes": [[1, 640, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7]]},
            {"dim": 1, "shapes": [[1, 640, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7]]},
            {"dim": 1, "shapes": [[1, 640, 7, 7], [1, 32, 7, 7], [1, 32, 7, 7]]},
            {"dim": 1, "shapes": [[1, 640, 7, 7], [1, 32, 7, 7]]},
            {"dim": 1, "shapes": [[1, 640, 7, 7]]},
            {"dim": -1, "shapes": [[1, 7, 7, 384], [1, 7, 7, 384], [1, 7, 7, 384], [1, 7, 7, 384]]},
            {"dim": -1, "shapes": [[1, 7, 7, 512], [1, 7, 7, 512], [1, 7, 7, 512], [1, 7, 7, 512]]},
            {"dim": -1, "shapes": [[1, 71, 7, 32], [1, 71, 7, 32]]},
            {"dim": 1, "shapes": [[150, 2], [150, 2]]},
            {"dim": 1, "shapes": [[2166, 2], [2166, 2]]},
            {"dim": 1, "shapes": [[24, 2], [24, 2]]},
            {"dim": 1, "shapes": [[2400, 2], [2400, 2]]},
            {"dim": 1, "shapes": [[6, 2], [6, 2]]},
            {"dim": 1, "shapes": [[600, 2], [600, 2]]},
            {"dim": -1, "shapes": [[8732, 2], [8732, 2]]},
        ],
        "dtype": [ttnn.bfloat16, ttnn.uint32, ttnn.int32],
        "layout": [ttnn.ROW_MAJOR_LAYOUT, ttnn.TILE_LAYOUT],
    }
}


# Invalidate vector is called during the generation phase where each vector will be passed in.
# If invalidated, the vector will still be stored but will be skipped.
# Returns False, None if the vector is valid, and True, str with a reason for invalidation if it is invalid.
def invalidate_vector(test_vector) -> Tuple[bool, Optional[str]]:
    if test_vector["layout"] == ttnn.ROW_MAJOR_LAYOUT:
        if test_vector["dtype"] == ttnn.bfloat8_b:
            return True, "bfloat8_b not supported with ROW_MAJOR_LAYOUT"

    return False, None


torch.manual_seed(0)


def random_torch_tensor(dtype, shape):
    if dtype == ttnn.uint16:
        return torch.randint(0, 100, shape).to(torch.int16)
    if dtype == ttnn.int32:
        return torch.randint(-(2**31), 2**31, shape, dtype=torch.int32)
    if dtype == ttnn.uint32:
        return torch.randint(0, 2**31, shape, dtype=torch.int32)
    return torch.rand(shape).bfloat16().float()


def run(
    concat_specs,
    dtype,
    layout,
    *,
    device,
) -> list:
    torch_input_tensors = [random_torch_tensor(dtype, shape) for shape in concat_specs["shapes"]]
    torch_output_tensor = torch.concat(torch_input_tensors, dim=concat_specs["dim"])

    ttnn_input_tensors = [
        ttnn.from_torch(torch_input_tensor, device=device, layout=layout, dtype=dtype)
        for torch_input_tensor in torch_input_tensors
    ]
    start_time = start_measuring_time()
    result_tensor = ttnn.concat(ttnn_input_tensors, dim=concat_specs["dim"])
    e2e_perf = stop_measuring_time(start_time)
    output_tensor = ttnn.to_torch(result_tensor)

    return [check_with_pcc(torch_output_tensor, output_tensor, 0.999), e2e_perf]


@pytest.mark.parametrize("concat_spec", parameters["nightly"]["concat_specs"])
@pytest.mark.parametrize("dtype", parameters["nightly"]["dtype"])
@pytest.mark.parametrize("layout", parameters["nightly"]["layout"])
def test_concat_pytorch2(concat_spec, dtype, layout, device):
    shapes = concat_spec["shapes"]
    dim = concat_spec["dim"]
    if dtype == ttnn.bfloat16 and any([shape[-1] % 2 != 0 for shape in shapes]) and layout == ttnn.ROW_MAJOR_LAYOUT:
        pytest.skip("Skipping test for RM bfloat16 with odd last dimension")

    torch_input_tensors = [random_torch_tensor(dtype, shape) for shape in shapes]
    torch_output_tensor = torch.cat(torch_input_tensors, dim=dim)

    ttnn_input_tensors = [
        ttnn.from_torch(torch_input_tensor, device=device, layout=layout, dtype=dtype)
        for torch_input_tensor in torch_input_tensors
    ]

    result_tensor = ttnn.concat(ttnn_input_tensors, dim=dim)
    output_tensor = ttnn.to_torch(result_tensor)

    assert check_with_pcc(
        torch_output_tensor, output_tensor, 0.999
    ), "Output tensors do not match within the specified precision"
